package org.liaohailong.library.util;

import android.content.Context;
import android.graphics.Bitmap;
import android.graphics.ImageFormat;
import android.media.Image;
import android.media.ImageReader;
import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaExtractor;
import android.media.MediaFormat;
import android.os.Handler;
import android.os.HandlerThread;
import android.util.Log;

import androidx.annotation.NonNull;
import androidx.annotation.WorkerThread;

import java.nio.ByteBuffer;
import java.util.LinkedList;

/**
 * Author: liaohailong
 * Time: 2020/10/29 14:06
 * Describe: 获取视频缩略图
 */
public class ThumbHelper {
    private static final String TAG = "ThumbHelper";

    private final Context context;
    private final LinkedList<ThumbInfo> cache = new LinkedList<>();

    public ThumbHelper(Context context) {
        this.context = context;
    }


    @WorkerThread
    public VideoInfo getVideoInfo(String filePath) {
        MediaExtractor extractor = null;
        try {
            extractor = new MediaExtractor();
            extractor.setDataSource(filePath);
            int trackCount = extractor.getTrackCount();
            // 找出并使用video轨道数据
            MediaFormat videoFormat = null;
            for (int i = 0; i < trackCount; i++) {
                MediaFormat trackFormat = extractor.getTrackFormat(i);
                String mime = trackFormat.getString(MediaFormat.KEY_MIME);
                if (mime != null && mime.contains("video")) {
                    videoFormat = trackFormat;
                    extractor.selectTrack(i);
                    break;
                }
            }

            if (videoFormat == null) {
                return null;
            }

            int formatWidth = videoFormat.getInteger(MediaFormat.KEY_WIDTH);
            int formatHeight = videoFormat.getInteger(MediaFormat.KEY_HEIGHT);
            String mimeType = videoFormat.getString(MediaFormat.KEY_MIME);
            long durationUs = videoFormat.getLong(MediaFormat.KEY_DURATION);

            return new VideoInfo(formatWidth, formatHeight, durationUs, mimeType);
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        } finally {
            if (extractor != null) {
                extractor.release();
            }
        }
    }

    /**
     * 开始解析视频缩略图
     *
     * @param filePath     视频本地路径
     * @param inSampleSize 缩小大小，类似BitmapFactory.Option.inSampleSize，值越大，图片越小，传入2的倍数
     * @param start        开始时长 单位：ms
     * @param interval     间隔时长 单位：ms
     * @return 任务标记，可以调用cancel函数，中止解析
     */
    @WorkerThread
    public TaskTag start(
            String filePath,
            int inSampleSize,
            int start,
            int interval,
            OnThumbReadCallback callback) {
        TaskTag taskTag = new TaskTag();
        MediaExtractor extractor = null;
        MediaCodec codec = null;
        try {
            extractor = new MediaExtractor();
            extractor.setDataSource(filePath);
            int trackCount = extractor.getTrackCount();
            // 找出并使用video轨道数据
            MediaFormat videoFormat = null;
            for (int i = 0; i < trackCount; i++) {
                MediaFormat trackFormat = extractor.getTrackFormat(i);
                String mime = trackFormat.getString(MediaFormat.KEY_MIME);
                if (mime != null && mime.contains("video")) {
                    videoFormat = trackFormat;
                    extractor.selectTrack(i);
                    break;
                }
            }

            // 文件里面没有视频轨道
            if (videoFormat == null) {
                return null;
            }

            int imageFormat = ImageFormat.YUV_420_888;
            int colorFormat = MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Flexible;
            int formatWidth = videoFormat.getInteger(MediaFormat.KEY_WIDTH);
            int formatHeight = videoFormat.getInteger(MediaFormat.KEY_HEIGHT);
            String mimeType = videoFormat.getString(MediaFormat.KEY_MIME);
            long durationUs = videoFormat.getLong(MediaFormat.KEY_DURATION);

            // 重新设置视频数据解析参数
            videoFormat.setInteger(MediaFormat.KEY_COLOR_FORMAT, colorFormat);
            videoFormat.setInteger(MediaFormat.KEY_WIDTH, formatWidth);
            videoFormat.setInteger(MediaFormat.KEY_HEIGHT, formatHeight);

            // 创建codec
            codec = MediaCodec.createDecoderByType(mimeType);
            // 创建image reader 接受图像数据
            ImageReader imageReader = ImageReader.newInstance(
                    formatWidth,
                    formatHeight,
                    imageFormat,
                    1);

            // 设置image reader 图像处理handler - 尽可能在子线程处理图像操作
            HandlerThread handlerThread = new HandlerThread("image-reader");
            handlerThread.start();
            Handler handler = new Handler(handlerThread.getLooper());
            imageReader.setOnImageAvailableListener(new MyOnImageAvailableListener(this, callback, inSampleSize), handler);

            // 配置codec的视频解码类型，和画面渲染surface
            codec.configure(videoFormat, imageReader.getSurface(), null, 0);
            codec.start();

            // 开始解码
            MediaCodec.BufferInfo bufferInfo = new MediaCodec.BufferInfo();
            long timeOut = 1000 * 5;
            boolean inputDone = false;
            boolean outputDone = false;
            int seekEndToCount = 0;
            int seekToCount = 0;
            int renderCount = 0;

            // 读取前先移动到指定起始帧
            long startUs = start * 1000;
            extractor.seekTo(startUs, MediaExtractor.SEEK_TO_PREVIOUS_SYNC);

            while (!outputDone) {
                if (taskTag.canceled) {
                    // 用户外部取消任务
                    break;
                }

                if (!inputDone) {
                    // 给codec填充视频数据
                    int inputBufferIndex = codec.dequeueInputBuffer(timeOut);
                    if (inputBufferIndex >= 0) {
                        ByteBuffer inputBuffer = codec.getInputBuffer(inputBufferIndex);

                        int sampleData = extractor.readSampleData(inputBuffer, 0);
                        long sampleTime = extractor.getSampleTime();
                        if (sampleData > 0) {
                            if (seekEndToCount >= 2) {
                                codec.queueInputBuffer(inputBufferIndex,
                                        0,
                                        0,
                                        0L,
                                        MediaCodec.BUFFER_FLAG_END_OF_STREAM);
                                inputDone = true;
                            } else
                                codec.queueInputBuffer(inputBufferIndex, 0, sampleData, sampleTime, 0);
                            // 继续解析
                            seekToCount++;
                            long seekToUs = startUs + (interval * 1000L) * seekToCount;
                            if (seekToUs >= durationUs) {
                                seekEndToCount++;
                            }
                            extractor.seekTo(seekToUs, MediaExtractor.SEEK_TO_PREVIOUS_SYNC);
                        } else {
                            // 小于0，说明读完了
                            codec.queueInputBuffer(inputBufferIndex,
                                    0,
                                    0,
                                    0L,
                                    MediaCodec.BUFFER_FLAG_END_OF_STREAM);
                            inputDone = true;
                        }
                    }
                }

                // 获取解码好的数据 status 有可能是buffer index 也有可能是状态flag
                int status = codec.dequeueOutputBuffer(bufferInfo, timeOut);
                switch (status) {
                    case MediaCodec.INFO_TRY_AGAIN_LATER:
                        // 继续
                        break;
                    case MediaCodec.INFO_OUTPUT_FORMAT_CHANGED:
                        // 开始进行解码
                        break;
                    case MediaCodec.INFO_OUTPUT_BUFFERS_CHANGED:
                        // 正在解码
                        break;
                    default:
                        // 判断当前编解码器的状态 - 是否已经解码完毕
                        if ((bufferInfo.flags & MediaCodec.BUFFER_FLAG_END_OF_STREAM) != 0) {
                            outputDone = true;
                            callback.onThumbReadComplete();
                        }
                        boolean doRender = (bufferInfo.size != 0);

                        // 记录渲染数量
                        if (doRender) {
                            cache.addLast(new ThumbInfo(bufferInfo.presentationTimeUs));
                            renderCount++;
                        }

                        // 把正在解析出来的截图数量打印一下
                        if (!outputDone && doRender)
                            Log.i(TAG, "presentationTimeUs = " + bufferInfo.presentationTimeUs + ", renderCount = " + renderCount);

                        // 直接送去显示
                        Thread.sleep(16);
                        codec.releaseOutputBuffer(status, doRender);
                        break;
                }
            }
        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (codec != null) {
                codec.stop();
                codec.release();
            }
            if (extractor != null) {
                extractor.release();
            }
        }
        return taskTag;
    }


    /**
     * 获取ImageReader获取视频画面数据
     */
    private static final class MyOnImageAvailableListener implements ImageReader.OnImageAvailableListener {
        private final ThumbHelper thumbHelper;
        private final OnThumbReadCallback callback;
        private final int inSampleSize;
        private final YuvToRgbConverter yuvToRgbConverter;

        public MyOnImageAvailableListener(ThumbHelper thumbHelper,
                                          OnThumbReadCallback callback,
                                          int inSampleSize) {
            this.thumbHelper = thumbHelper;
            this.callback = callback;
            this.inSampleSize = inSampleSize;
            yuvToRgbConverter = new YuvToRgbConverter(thumbHelper.context);
        }

        @Override
        public void onImageAvailable(ImageReader reader) {
            try (Image img = reader.acquireLatestImage()) {
                if (img != null) {
                    // 这里得到YUV数据，需要讲YUV数据变成Bitmap
                    Image.Plane[] planes = img.getPlanes();
                    if (planes[0].getBuffer() == null) {
                        return;
                    }

                    long current = System.currentTimeMillis();
                    Bitmap bitmap = getBitmap(img, inSampleSize);
                    long cost = System.currentTimeMillis() - current;
                    Log.i(TAG, "getBitmap with inSampleSize = " + inSampleSize + ", cost: " + cost + "ms");
                    if (callback != null) {
                        ThumbInfo thumbInfo = thumbHelper.cache.pollFirst();
                        if (thumbInfo == null) thumbInfo = new ThumbInfo(-1);
                        thumbInfo.bitmap = bitmap;
                        callback.onThumbRead(thumbInfo);
                    }
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        }

        private Bitmap getBitmap(Image image, int inSampleSize) {
            int width = image.getWidth();
            int height = image.getHeight();

            Bitmap bitmap = Bitmap.createBitmap(width, height, Bitmap.Config.ARGB_8888);
            yuvToRgbConverter.yuvToRgb(image, bitmap);
            int dstWidth = width / inSampleSize;
            int dstHeight = height / inSampleSize;
            Bitmap result = Bitmap.createScaledBitmap(bitmap, dstWidth, dstHeight, false);
            if (!bitmap.isRecycled()) bitmap.recycle();
            return result;
        }

    }

    public interface OnThumbReadCallback {
        /**
         * @param thumbInfo 解析出来的每一帧画面
         */
        void onThumbRead(@NonNull ThumbInfo thumbInfo);

        /**
         * 视频全部解析完毕
         */
        void onThumbReadComplete();
    }

    public static final class ThumbInfo {
        private final long timeUs;
        private Bitmap bitmap;

        public ThumbInfo(long timeUs) {
            this.timeUs = timeUs;
        }

        public Bitmap getBitmap() {
            return bitmap;
        }

        public long getTimeUs() {
            return timeUs;
        }

        public long getTimeMs() {
            return getTimeUs() / 1000;
        }

        @Override
        public String toString() {
            return "ThumbInfo{" +
                    "timeUs=" + timeUs +
                    '}';
        }
    }

    public static final class TaskTag {
        private boolean canceled = false;

        public void cancel() {
            this.canceled = true;
        }
    }

    public static final class VideoInfo {
        private final int width;
        private final int height;
        private final long durationUs;
        private final String mime;

        public VideoInfo(int width, int height, long durationUs, String mime) {
            this.width = width;
            this.height = height;
            this.durationUs = durationUs;
            this.mime = mime;
        }

        public int getWidth() {
            return width;
        }

        public int getHeight() {
            return height;
        }

        public long getDurationUs() {
            return durationUs;
        }

        public long getDurationMs() {
            return getDurationUs() / 1000;
        }

        public String getMime() {
            return mime;
        }
    }
}
